async_utils.py 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. import inspect
  2. import typing as t
  3. from functools import WRAPPER_ASSIGNMENTS
  4. from functools import wraps
  5. from .utils import _PassArg
  6. from .utils import pass_eval_context
  7. V = t.TypeVar("V")
  8. def async_variant(normal_func): # type: ignore
  9. def decorator(async_func): # type: ignore
  10. pass_arg = _PassArg.from_obj(normal_func)
  11. need_eval_context = pass_arg is None
  12. if pass_arg is _PassArg.environment:
  13. def is_async(args: t.Any) -> bool:
  14. return t.cast(bool, args[0].is_async)
  15. else:
  16. def is_async(args: t.Any) -> bool:
  17. return t.cast(bool, args[0].environment.is_async)
  18. # Take the doc and annotations from the sync function, but the
  19. # name from the async function. Pallets-Sphinx-Themes
  20. # build_function_directive expects __wrapped__ to point to the
  21. # sync function.
  22. async_func_attrs = ("__module__", "__name__", "__qualname__")
  23. normal_func_attrs = tuple(set(WRAPPER_ASSIGNMENTS).difference(async_func_attrs))
  24. @wraps(normal_func, assigned=normal_func_attrs)
  25. @wraps(async_func, assigned=async_func_attrs, updated=())
  26. def wrapper(*args, **kwargs): # type: ignore
  27. b = is_async(args)
  28. if need_eval_context:
  29. args = args[1:]
  30. if b:
  31. return async_func(*args, **kwargs)
  32. return normal_func(*args, **kwargs)
  33. if need_eval_context:
  34. wrapper = pass_eval_context(wrapper)
  35. wrapper.jinja_async_variant = True
  36. return wrapper
  37. return decorator
  38. _common_primitives = {int, float, bool, str, list, dict, tuple, type(None)}
  39. async def auto_await(value: t.Union[t.Awaitable["V"], "V"]) -> "V":
  40. # Avoid a costly call to isawaitable
  41. if type(value) in _common_primitives:
  42. return t.cast("V", value)
  43. if inspect.isawaitable(value):
  44. return await t.cast("t.Awaitable[V]", value)
  45. return t.cast("V", value)
  46. async def auto_aiter(
  47. iterable: "t.Union[t.AsyncIterable[V], t.Iterable[V]]",
  48. ) -> "t.AsyncIterator[V]":
  49. if hasattr(iterable, "__aiter__"):
  50. async for item in t.cast("t.AsyncIterable[V]", iterable):
  51. yield item
  52. else:
  53. for item in t.cast("t.Iterable[V]", iterable):
  54. yield item
  55. async def auto_to_list(
  56. value: "t.Union[t.AsyncIterable[V], t.Iterable[V]]",
  57. ) -> t.List["V"]:
  58. return [x async for x in auto_aiter(value)]